import os

import numpy as np
from fpdf import FPDF

from utils import ModelSaver, gen_softmax


# 整理检查点的策略，并打印成PDF


def get_policy(path):
    saver = ModelSaver()
    pp1, pp2 = saver.load(path)

    nums = 2  # 概率要保留的小数点位数
    record = []

    for p in pp1:
        v1 = gen_softmax(pp1[p])
        v2 = gen_softmax(pp2[p])

        # 去除无效状态动作对
        if np.abs(v2 - v2[0]).sum() == 0 and np.abs(v1 - v1[0]).sum() == 0:
            continue

        line = str(p) + ":  "

        for i in np.round(v1, nums):
            s = str(i)
            s += "0" * (nums + 2 - len(s))
            line += s + "  "
        line += " |   "
        for i in np.round(v2, nums):
            s = str(i)
            s += "0" * (nums + 2 - len(s))
            line += s + "  "
        record.append(line)
    return record


def write_pdf(record, position):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font('Times', '', 14)  # 必须要， 字体，类型（粗体B，斜体I）， 大小
    for line in record:
        # 第一个参数 cell宽度 左距；第二个参数 cell高度
        pdf.cell(40, 10, txt=line, ln=1)  # ln 该单位格后的光标位置 1：下行开始
        # pdf.ln()        # 换行
    pdf.output(position, 'F')


class Evaluator():
    def __init__(self):
        from direct.spiel_env import Env
        self.env = Env(6, 1)
        self.saver = ModelSaver()  # 用于加载权重

    def reset_checkpoint(self, path):
        self.policy1, self.policy2 = self.saver.load(path)

    def get_actions(self, obs):
        softmax1 = gen_softmax(self.policy1[obs])
        action1 = np.random.choice(range(len(softmax1)), p=softmax1)  # 动作0合法， action即index

        softmax2 = gen_softmax(self.policy2[obs])
        action2 = np.random.choice(range(len(softmax2)), p=softmax2)

        return action1, action2, np.round(softmax1, 2), np.round(softmax2, 2)

    def sample_one_eps(self):
        lines = []
        obs, done = self.env.init()
        while not done:
            line = str(obs)
            action1, action2, softmax1, softmax2 = self.get_actions(obs)
            line += f", {action1}, {action2}:  "
            for i in softmax1:
                line += f"{str(i)}  "
            line += "  |    "
            for i in softmax1:
                line += f"{str(i)}  "
            obs, done = self.env.step((action1, action2))

            lines.append(line)
        return lines


def display_checkpoint(checkpoint_path, train_dir):
    record = get_policy(checkpoint_path)
    position = checkpoint_path.split(train_dir)[0] + "checkpoint_policy/"
    if not os.path.exists(position):
        os.makedirs(position)
    position += train_dir + "_" + os.path.basename(checkpoint_path).replace("pkl", "pdf")

    write_pdf(record, position)


def display_checkpoints(checkpoint_dir):

    paths = os.listdir(checkpoint_dir)
    paths = [str(j) for j in sorted([int(i.split(".")[0]) for i in paths])]
    record = []
    evaluator = Evaluator()
    for index, path in enumerate(paths):
        record.append("*" * 20 + "update times: " + path + "*" * 20)
        evaluator.reset_checkpoint(checkpoint_dir + f"/{path}.pkl")
        info = evaluator.sample_one_eps()
        record.extend(info)
    position = checkpoint_dir.split(r"/checkpoints")[0] + ".pdf"
    write_pdf(record, position)


def remote():
    checkpoints = [20000, 30000, 40000, 50000]
    for train_dir in os.listdir("/root/projects/jzhou/Alesia/0518/direct/train_info"):
        if ".pdf" in train_dir or "checkpoint" in train_dir:
            continue
        checkpoint_dir = rf"/root/projects/jzhou/Alesia/0518/direct/train_info/{train_dir}/checkpoints"

        # display_checkpoints(checkpoint_dir)

        for checkpoint in checkpoints:
            checkpoint_path = fr"{checkpoint_dir}/{checkpoint}.pkl"
            display_checkpoint(checkpoint_path, train_dir)


def local():
    checkpoints = [30000, 40000, 50000]

    checkpoint_dir = rf"E:\CodeHub\实验室项目\alesia\Alesia\analysis\train15\train_info\eg_0.9_-_solid_direct_No0\checkpoints"
    train_dir = checkpoint_dir.split("\\")[-2]
    for checkpoint in checkpoints:
        checkpoint_path = fr"{checkpoint_dir}/{checkpoint}.pkl"
        display_checkpoint(checkpoint_path, train_dir)


if __name__ == '__main__':
    local()
